import torch


class Conv(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0)):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding,
                                    bias=False)
        self.bn = torch.nn.BatchNorm2d(out_channels)
        self.relu = torch.nn.LeakyReLU()

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))


class Residual(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv_1 = Conv(in_channels, out_channels, (1, 1))
        self.conv_2 = Conv(out_channels, in_channels, kernel_size=(3, 3), padding=(1, 1))

    def forward(self, x):
        x1 = self.conv_2(self.conv_1(x))
        return x + x1


class ResBlock(torch.nn.Module):
    def __init__(self, n, in_channels):
        super().__init__()

        out_channels = in_channels * 2
        self.conv = Conv(in_channels, out_channels, (3, 3), stride=(2, 2), padding=(1, 1))

        self.residual_blocks = torch.nn.Sequential()

        for i in range(n):
            self.residual_blocks.add_module(str(i), Residual(out_channels, in_channels))

    def forward(self, x):
        x1 = self.conv(x)
        x2 = self.residual_blocks(x1)
        return x2


class YOLOBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.yolo_block = torch.nn.Sequential(
            Conv(in_channels, out_channels, (1, 1)),

            Conv(in_channels=out_channels, out_channels=2 * out_channels, kernel_size=(3, 3),
                            padding=(1, 1)),

            Conv(2 * out_channels, out_channels, (1, 1)),

            Conv(out_channels, 2 * out_channels, (3, 3), padding=(1, 1)),

            Conv(2 * out_channels, out_channels, kernel_size=(1, 1))

        )

    def forward(self, x):
        x1 = self.yolo_block(x)
        return x1


class YOLOModel(torch.nn.Module):

    def __init__(self, n):
        super().__init__()

        self.backbone_sequential_1 = torch.nn.Sequential(

            Conv(3, 32, (3, 3), padding=(1, 1)),  # 32,256,256
            ResBlock(1, 32),  # 64,128,128
            ResBlock(2, 64),  # 128,64,64
            ResBlock(8, 128)  # 256,32,32
        )

        # ----------------

        self.backbone_residual_4 = ResBlock(8, 256)  # 512,16,16

        self.backbone_sequential_2 = torch.nn.Sequential(
            ResBlock(4, 512),  # 1024,8,8
        )

        self.yolo_block = YOLOBlock(1024, 512)  # 512,8,8

        self.sequential_6 = torch.nn.Sequential(
            Conv(512, 1024, (3, 3), padding=(1, 1)),
            Conv(1024, n, (1, 1))
        )

        # ------------------

        self.conv_2 = Conv(512, 1024, (3, 3), padding=(1, 1))  # 1024,8,8

        self.conv_3 = Conv(1024, n, (1, 1))  # n,8,8

        self.sequential_2 = torch.nn.Sequential(
            Conv(512, 256, (1, 1)),
            torch.nn.Upsample(scale_factor=(2, 2)),

        )

        self.yolo_block_1 = YOLOBlock(768, 256)

        self.sequential_3 = torch.nn.Sequential(
            Conv(256, 128, (1, 1)),
            torch.nn.Upsample(scale_factor=(2, 2))
        )

        self.sequential_4 = torch.nn.Sequential(
            YOLOBlock(384, 128),
            Conv(128, 256, (3, 3), padding=(1, 1)),
            Conv(256, n, (1, 1))
        )

        self.sequential_5 = torch.nn.Sequential(
            Conv(256, 512, (3, 3), padding=(1, 1)),
            Conv(512, n, (1, 1))
        )

    def forward(self, x):
        x1 = self.backbone_sequential_1(x)
        x2 = self.backbone_residual_4(x1)

        x3 = self.yolo_block(self.backbone_sequential_2(x2))

        x4 = self.yolo_block_1(torch.cat([x2, self.sequential_2(x3)], dim=1))

        large_x = self.sequential_4(torch.cat([x1, self.sequential_3(x4)], dim=1))
        medium_x = self.sequential_5(x4)
        min_x = self.sequential_6(x3)

        return min_x, medium_x, large_x
